#!/bin/bash

# 模拟系统性能瓶颈脚本
# 功能：模拟CPU、内存、磁盘和网络性能瓶颈，用于测试性能监控脚本
# 作者：Xiao Ru An
# 日期：$(date +"%Y-%m-%d")

# 设置颜色
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m' # 无颜色

# 打印信息
print_info() {
    echo -e "${GREEN}$1${NC}"
}

print_warning() {
    echo -e "${YELLOW}$1${NC}"
}

print_error() {
    echo -e "${RED}$1${NC}"
}

# 显示帮助信息
show_help() {
    echo "使用方法: $0 [选项]"
    echo "选项:"
    echo "  -c, --cpu       模拟CPU性能瓶颈"
    echo "  -m, --memory    模拟内存性能瓶颈"
    echo "  -d, --disk      模拟磁盘性能瓶颈"
    echo "  -n, --network   模拟网络性能瓶颈"
    echo "  -s, --stop      停止所有模拟的瓶颈"
    echo "  -h, --help      显示此帮助信息"
    echo ""
    echo "示例:"
    echo "  $0 --cpu        模拟CPU瓶颈"
    echo "  $0 -m           模拟内存瓶颈"
    echo "  $0 --stop       停止所有模拟"
}

# 模拟CPU瓶颈
simulate_cpu_bottleneck() {
    print_info "正在模拟CPU性能瓶颈..."
    print_info "启动10个CPU密集型进程，每个进程占用一个CPU核心的100%资源。"
    
    # 创建PID文件来跟踪进程
    PID_FILE="/tmp/cpu_bottleneck_pids.txt"
    > "$PID_FILE"  # 清空文件
    
    # 启动CPU密集型进程
    for i in {1..10}; do
        (while true; do echo "scale=1000; 4*a(1)" | bc -l > /dev/null; done) &
        echo $! >> "$PID_FILE"
    done
    
    print_info "CPU瓶颈模拟已启动。PID信息保存在 $PID_FILE"
    print_warning "要停止模拟，请运行: $0 --stop"
}

# 模拟内存瓶颈
simulate_memory_bottleneck() {
    print_info "正在模拟内存性能瓶颈..."
    print_info "分配大量内存并保持占用。"
    
    # 创建PID文件来跟踪进程
    PID_FILE="/tmp/memory_bottleneck_pids.txt"
    > "$PID_FILE"  # 清空文件
    
    # 启动内存占用进程
    (python3 -c "import time;a=[];a.append('A'*1024*1024*100);time.sleep(3600)" 2>/dev/null || \
     python -c "import time;a=[];a.append('A'*1024*1024*100);time.sleep(3600)" 2>/dev/null || \
     perl -e "\$a='A' x (1024*1024*100);sleep(3600)" 2>/dev/null) &
    
    echo $! >> "$PID_FILE"
    
    print_info "内存瓶颈模拟已启动。PID信息保存在 $PID_FILE"
    print_warning "要停止模拟，请运行: $0 --stop"
}

# 模拟磁盘瓶颈
simulate_disk_bottleneck() {
    print_info "正在模拟磁盘性能瓶颈..."
    print_info "进行大量磁盘读写操作。"
    
    # 创建PID文件来跟踪进程
    PID_FILE="/tmp/disk_bottleneck_pids.txt"
    > "$PID_FILE"  # 清空文件
    
    # 创建临时目录
    TEMP_DIR="/tmp/disk_bottleneck"
    mkdir -p "$TEMP_DIR"
    
    # 启动磁盘I/O密集型进程
    (while true; do
        dd if=/dev/zero of=$TEMP_DIR/testfile bs=1M count=100 2>/dev/null
        dd if=$TEMP_DIR/testfile of=/dev/null bs=1M 2>/dev/null
        rm -f $TEMP_DIR/testfile
    done) &
    
    echo $! >> "$PID_FILE"
    
    print_info "磁盘瓶颈模拟已启动。PID信息保存在 $PID_FILE"
    print_warning "要停止模拟，请运行: $0 --stop"
}

# 模拟网络瓶颈
simulate_network_bottleneck() {
    print_info "正在模拟网络性能瓶颈..."
    print_info "进行大量网络连接和数据传输。"
    
    # 创建PID文件来跟踪进程
    PID_FILE="/tmp/network_bottleneck_pids.txt"
    > "$PID_FILE"  # 清空文件
    
    # 启动网络密集型进程
    # 如果有wget或curl，使用它们下载大文件
    if command -v wget >/dev/null 2>&1; then
        (while true; do
            wget -O /dev/null https://speed.hetzner.de/100MB.bin 2>/dev/null
        done) &
        echo $! >> "$PID_FILE"
    elif command -v curl >/dev/null 2>&1; then
        (while true; do
            curl -o /dev/null https://speed.hetzner.de/100MB.bin 2>/dev/null
        done) &
        echo $! >> "$PID_FILE"
    else
        # 如果没有wget或curl，则使用nc创建大量连接
        (while true; do
            (for i in {1..50}; do
                nc -z google.com 80 2>/dev/null &
            done
            sleep 5) &
        done) &
        echo $! >> "$PID_FILE"
    fi
    
    print_info "网络瓶颈模拟已启动。PID信息保存在 $PID_FILE"
    print_warning "要停止模拟，请运行: $0 --stop"
}

# 停止所有模拟
stop_all_bottlenecks() {
    print_info "正在停止所有模拟的性能瓶颈..."
    
    # 停止CPU瓶颈进程
    if [ -f "/tmp/cpu_bottleneck_pids.txt" ]; then
        while read pid; do
            if kill -0 $pid 2>/dev/null; then
                kill $pid
                print_info "已停止CPU瓶颈进程: $pid"
            fi
        done < "/tmp/cpu_bottleneck_pids.txt"
        rm -f "/tmp/cpu_bottleneck_pids.txt"
    fi
    
    # 停止内存瓶颈进程
    if [ -f "/tmp/memory_bottleneck_pids.txt" ]; then
        while read pid; do
            if kill -0 $pid 2>/dev/null; then
                kill $pid
                print_info "已停止内存瓶颈进程: $pid"
            fi
        done < "/tmp/memory_bottleneck_pids.txt"
        rm -f "/tmp/memory_bottleneck_pids.txt"
    fi
    
    # 停止磁盘瓶颈进程
    if [ -f "/tmp/disk_bottleneck_pids.txt" ]; then
        while read pid; do
            if kill -0 $pid 2>/dev/null; then
                kill $pid
                print_info "已停止磁盘瓶颈进程: $pid"
            fi
        done < "/tmp/disk_bottleneck_pids.txt"
        rm -f "/tmp/disk_bottleneck_pids.txt"
    fi
    
    # 停止网络瓶颈进程
    if [ -f "/tmp/network_bottleneck_pids.txt" ]; then
        while read pid; do
            if kill -0 $pid 2>/dev/null; then
                kill $pid
                print_info "已停止网络瓶颈进程: $pid"
            fi
        done < "/tmp/network_bottleneck_pids.txt"
        rm -f "/tmp/network_bottleneck_pids.txt"
    fi
    
    # 清理临时文件
    rm -rf "/tmp/disk_bottleneck"
    
    print_info "所有模拟的性能瓶颈已停止。"
}

# 主函数
main() {
    # 检查参数
    if [ $# -eq 0 ]; then
        show_help
        exit 1
    fi
    
    case "$1" in
        -c|--cpu)
            simulate_cpu_bottleneck
            ;;
        -m|--memory)
            simulate_memory_bottleneck
            ;;
        -d|--disk)
            simulate_disk_bottleneck
            ;;
        -n|--network)
            simulate_network_bottleneck
            ;;
        -s|--stop)
            stop_all_bottlenecks
            ;;
        -h|--help)
            show_help
            ;;
        *)
            print_error "未知选项: $1"
            show_help
            exit 1
            ;;
    esac
}

# 执行主函数
main "$@"